package org.btik.lightweb.sensitive;

import com.google.gson.Gson;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpSession;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.btik.light.server.platform.common.api.bean.sensitive.annotation.SensitiveBody;
import org.btik.light.server.platform.common.api.bean.sensitive.annotation.SensitiveField;
import org.btik.light.server.platform.common.api.bean.sensitive.annotation.SensitiveParams;
import org.btik.light.server.platform.common.api.bean.sensitive.annotation.StringEncoding;
import org.btik.light.server.platform.common.api.service.SensitiveDataService;
import org.btik.light.tool.AESUtil;
import org.btik.light.tool.ByteUtil;
import org.btik.light.tool.DhUtil;
import org.btik.light.tool.type.StringUtil;
import org.btik.light.tool.type.TypeUtil;
import org.btik.lightweb.client.login.LoginController;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.support.WebDataBinderFactory;
import org.springframework.web.context.request.NativeWebRequest;
import org.springframework.web.method.support.HandlerMethodArgumentResolver;
import org.springframework.web.method.support.ModelAndViewContainer;

import java.beans.PropertyDescriptor;
import java.io.IOException;
import java.io.InputStreamReader;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.security.InvalidKeyException;
import java.util.Arrays;
import java.util.Base64;
import java.util.Objects;

/**
 * @author lustre
 * @since 2023/6/8 10:36
 */
@Component
public class SensitiveDtoResolver implements HandlerMethodArgumentResolver {

    private final Log log = LogFactory.getLog(SensitiveDtoResolver.class);

    private static final String EXCHANGE_SESSION_ID = "exchangeSessionId";

    public static final String DH_PUBLIC_KEY = "publicKey";

    private static final Gson gson = new Gson();

    @Autowired
    private SensitiveDataService sensitiveDataService;

    @Override
    public boolean supportsParameter(MethodParameter parameter) {
        return parameter.getParameter().isAnnotationPresent(SensitiveBody.class) || parameter.getParameter().isAnnotationPresent(SensitiveParams.class);
    }

    @Override
    public Object resolveArgument(MethodParameter parameter, ModelAndViewContainer mavContainer, NativeWebRequest webRequest, WebDataBinderFactory binderFactory) throws Exception {


        HttpServletRequest nativeRequest = webRequest.getNativeRequest(HttpServletRequest.class);
        if (nativeRequest == null) {
            throw new IllegalStateException("can resolve request " + webRequest);
        }

        byte[] secretKey = null;
        Parameter parameter1 = parameter.getParameter();
        Class<?> parameterType = parameter.getParameterType();
        Object o = null;
        if (parameter1.isAnnotationPresent(SensitiveBody.class)) {
            SensitiveBody sensitiveBody = parameter1.getAnnotation(SensitiveBody.class);
            boolean userSession = sensitiveBody.userSession();
            secretKey = exchangeFinal(nativeRequest, userSession);
            o = parseJson(parameterType, nativeRequest, secretKey);
        } else if (parameter1.isAnnotationPresent(SensitiveParams.class)) {
            SensitiveParams sensitiveParams = parameter1.getAnnotation(SensitiveParams.class);
            secretKey = exchangeFinal(nativeRequest, sensitiveParams.userSession());
            o = parseParams(parameter1, nativeRequest, secretKey, sensitiveParams);
        }

        if (secretKey != null) {
            Arrays.fill(secretKey, (byte) 0);
        }


        return o;
    }

    /**
     * 协商出密钥
     */
    private byte[] exchangeFinal(HttpServletRequest nativeRequest, boolean userSession) {

        HttpSession session = nativeRequest.getSession(false);

        String publicKey = nativeRequest.getHeader(DH_PUBLIC_KEY);
        if (userSession) {
            if(session == null){
                throw new IllegalStateException("session Illegal ");
            }
            Object dhKeyPair = session.getAttribute(LoginController.DH_KEY_PAIR);
            if (!(dhKeyPair instanceof DhUtil.DhKeyPair dhKeyPair1)) {
                throw new IllegalStateException("session Illegal ");
            }
            return sensitiveDataService.exchangeFinal(publicKey, dhKeyPair1);
        }

        String exchangeSessionId = nativeRequest.getHeader(EXCHANGE_SESSION_ID);
        if (StringUtil.isEmpty(exchangeSessionId)) {
            throw new IllegalStateException("ex session Illegal ");
        }

        return sensitiveDataService.exchangeFinal(publicKey, exchangeSessionId);
    }

    private Object parseParams(Parameter parameter, HttpServletRequest nativeRequest, byte[] secretKey, SensitiveParams annotation) throws InvalidKeyException {
        Class<?> type = parameter.getType();
        if (TypeUtil.isCommonType(type)) {

            String name = annotation.value();
            if (StringUtil.isEmpty(name)) {
                name = parameter.getName();
            }
            StringEncoding encoding = annotation.encoding();
            String nameValue = nativeRequest.getParameter(name);
            byte[] decrypt = decrypt(secretKey, nameValue, encoding);
            return TypeUtil.parseValueFromString(new String(decrypt), type);
        }


        return null;

    }

    private Object parseJson(Class<?> parameterType, HttpServletRequest nativeRequest, byte[] secretKey) throws IOException, InvocationTargetException, IllegalAccessException, InvalidKeyException {

        Object o = gson.fromJson(new InputStreamReader(nativeRequest.getInputStream()), parameterType);


        PropertyDescriptor[] propertyDescriptors = BeanUtils.getPropertyDescriptors(parameterType);
        for (PropertyDescriptor propertyDescriptor : propertyDescriptors) {
            Method writeMethod = propertyDescriptor.getWriteMethod();
            if (writeMethod == null || !writeMethod.isAnnotationPresent(SensitiveField.class)) {
                continue;
            }

            Method readMethod = propertyDescriptor.getReadMethod();
            Object invoke = readMethod.invoke(o);
            if (!(invoke instanceof String ciphertext)) {
                continue;
            }
            StringEncoding stringEncoding = writeMethod.getAnnotation(SensitiveField.class).value();
            String plaintext = new String(decrypt(secretKey, ciphertext, stringEncoding));
            writeMethod.invoke(o, plaintext);
        }


        return o;
    }

    private byte[] decrypt(byte[] secretKey, String ciphertext, StringEncoding stringEncoding) throws InvalidKeyException {
        byte[] ciphertextBytes;
        if (Objects.requireNonNull(stringEncoding) == StringEncoding.HEX_STRING) {
            ciphertextBytes = ByteUtil.hexStrToByteArray(ciphertext);
        } else {
            ciphertextBytes = Base64.getDecoder().decode(ciphertext);
        }

        byte[] decrypt = AESUtil.decrypt(AESUtil.ECB_PKCS5Padding,
                ciphertextBytes, secretKey);
        if (decrypt == null) {
            log.info("decrypt failed");
            throw new IllegalStateException("Illegal parameters");
        }
        return decrypt;
    }
}
